import numpy
import matplotlib.pyplot as plt


def export_ridge(clf):
  pass


def ridge_regress(xdf, y):
  from sklearn.linear_model import Ridge
  clf = Ridge(alpha=2, normalize=True)
  clf.fit(xdf, y)
  return clf


def feature_select(xdf, y):
  chosen = []
  corrs = []
  for xname in xdf.columns:
    corrs.append(abs(numpy.corrcoef(xdf[xname], y)[0, 1]))

  min_corr = 0.05
  if len(corrs) > 150:
    min_corr = max(min_corr, sorted(corrs, reverse=True)[150])

  for xname in xdf.columns:
    if not xname.startswith("X."):
      continue
    if abs(numpy.corrcoef(xdf[xname], y)[0, 1]) > min_corr:
      chosen.append(xname)
  return chosen


def _chunks(l, n):
  m = int(numpy.ceil(len(l) / n))
  for i in range(0, len(l), m):
    yield l[i:i + m]


def get_quantile(x_input, y_input, nchunks=30):
  try:
    x, _, y = zip(*sorted(zip(x_input, range(len(x_input)), y_input)))
    avg_ys = [numpy.mean(y_chunk) for y_chunk in _chunks(y, nchunks)]
    avg_xs = [numpy.mean(x_chunk) for x_chunk in _chunks(x, nchunks)]
    return avg_xs, avg_ys
  except Exception:
    return [], []


def plot_model(title, timestamp, midp, clf, Xin, yin, Xout, yout, yhoriz, filepath):
  import scipy.stats as ss
  yfit = clf.predict(Xin)
  ypred = clf.predict(Xout)
  xqi, yqi = get_quantile(yfit, yin)
  xqo, yqo = get_quantile(ypred, yout)
  plt.subplot(421)
  plt.plot(xqi, yqi, 'g.-', markersize=1, linewidth=0.5)
  plt.title("%s is  (# %d)" % (title, len(yin)))
  plt.subplot(422)
  plt.plot(xqo, yqo, 'g.-', markersize=1, linewidth=0.5)
  plt.title("%s oos (# %d)" % (title, len(yout)))
  plt.subplot(423)
  plt.plot(ypred, yout, 'g.', markersize=1, linewidth=0.5)
  plt.title("%s oos (# %d)" % (title, len(yout)))
  plt.subplot(424)
  plt.plot(yin, 'g.', yout, 'g.', markersize=1, linewidth=0.5)
  plt.subplot(425)
  plt.plot(midp, 'g-', markersize=1, linewidth=0.5)
  plt.subplot(426)
  plt.plot(yhoriz, 'g-', markersize=1, linewidth=0.5)
  plt.subplot(4, 3, 10)
  # plt.plot(sorted(ypred), 'g-', linewidth=0.5)
  plt.hist(ypred, 50)
  plt.title("y oos value dist")
  plt.subplot(4, 3, 11)
  plt.plot(ypred, 'g.', markersize=1, linewidth=0.5)
  plt.title("y oos value")
  plt.subplot(4, 3, 12)
  plt.plot((numpy.clip(ss.zscore(yfit), -2, 2) * yin).cumsum(),
           'g-', (numpy.clip(ss.zscore(ypred), -2, 2) * yout).cumsum(),
           'g-',
           markersize=1,
           linewidth=0.5)
  plt.savefig(filepath, dpi=200)
  plt.close()


def get_extra_info(clf, Xin, Xout, yin, yout):
  ypred = clf.predict(Xout)
  xqo, yqo = get_quantile(ypred, yout)
  return {
      "rsq_oos": numpy.corrcoef(ypred, yout)[0, 1],
      "predinfo_oos": {
          "ypred_qtl": xqo,
          "yact_qtl": yqo,
      }
  }
